"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：贝叶斯选币炼丹炉
"""

from datetime import datetime
import os
import shutil
import pandas as pd
from program import Config as Cfg
import program.plumesoft_bys.bys_function as func
from program.plumesoft_bys.bys_back_test import batch_run_stg_back_test
from program.plumesoft_bys.bys_core import Bayesian_Search
from program.plumesoft_bys.bys_common import init_env, log, dbg


def run_top_stg_back_test(
    params: dict,
    bys,
    df: pd.DataFrame,
    cls,
    name: str,
    disp_name: str,
    back_test_date_start,
    back_test_date_end,
    bys_back_test_eval_score_dict,
    update_score: bool,
):
    # 生成临时策略文件用于计算样本外表现，生成最终的评分
    stg_file_list = bys.output_result_stg_files(
        params, df, len(df), 0, "炼丹回测", open_pic=False
    )

    df[f"{name}回测数据"] = str({})
    df[f"{name}评分"] = 0
    if len(stg_file_list) == len(df):
        back_test_dt = {}
        log(f"开始对 {disp_name} 生成的策略文件进行样本外回测，计算最终的综合评分...")
        log(f"回测区间：{back_test_date_start} - {back_test_date_end}")
        ret_list = batch_run_stg_back_test(
            params["output_stg_folder"],
            stg_file_list,
            date_start=back_test_date_start,
            date_end=back_test_date_end,
            if_plot=False,
            n_jobs=params["stg_back_test_n_jobs"],
        )
        for stg_file, data_dt in ret_list:
            # 取得策略索引号
            stg_bt_dt = back_test_dt.get(stg_file, {})

            if set(list(data_dt.keys())) & set(
                list(bys_back_test_eval_score_dict.keys())
            ):
                # 评价表中有对应类型的评分
                eval_dt = bys_back_test_eval_score_dict
            else:
                # 兼容早期版本没有区分不同类型的评分
                eval_dt = {"中性": bys_back_test_eval_score_dict}

            # 计算样本外评分
            score_sum = stg_bt_dt.get(f"{name}评分", 0)
            core_dt = stg_bt_dt.get(f"{name}回测数据", {})
            for tp, data_list in data_dt.items():
                if tp in eval_dt.keys():
                    # 计算评分
                    score_select, r_dt, c_dt = func.cal_rtn_score(
                        params,
                        cls,
                        tp,
                        eval_dt[tp],
                        data_list[0],
                        data_list[1],
                        data_list[2],
                        data_list[3],
                        is_sample=False,
                    )
                    score_sum += score_select
                    core_dt[f"{tp}_回测结果"] = r_dt
                    core_dt[f"{tp}_评分明细"] = c_dt

            # 策略数据
            stg_bt_dt[f"{name}回测数据"] = core_dt
            stg_bt_dt[f"{name}评分"] = score_sum
            back_test_dt[stg_file] = stg_bt_dt

        # 更新样本外回测数据到最终结果
        for stg_file, stg_bt_dt in back_test_dt.items():
            idx = stg_file_list.index(stg_file)
            df.loc[idx, f"{name}回测数据"] = str(stg_bt_dt[f"{name}回测数据"])
            df.loc[idx, f"{name}评分"] = stg_bt_dt[f"{name}评分"]

        if update_score:
            # 原来的 score 改名为 炼丹评分
            df.rename(columns={"score": "炼丹评分"}, inplace=True)
            # 结合炼丹和样本外回测计算总评分
            df["score"] = (
                df["炼丹评分"] * params["study_batch_sample_score_pwr"]
                + df[f"{name}评分"] * params["study_batch_back_test_score_pwr"]
            )
    else:
        log(f"保存临时策略文件数量 {len(stg_file_list)} != {len(df)}，无法进行回测")
        exit()

    log("删除用来回测的临时策略文件...")
    for stg_file in stg_file_list:
        os.remove(
            os.path.join(
                Cfg.root_path,
                "program",
                params["output_stg_folder"],
                f"{stg_file}.py",
            )
        )
    return df


def get_back_test_eval_score_dict(params):
    # 获取样本外表现的评价字典
    bys_back_test_eval_score_dict = params["bys_back_test_eval_score_dict"]
    if bys_back_test_eval_score_dict is None or len(bys_back_test_eval_score_dict) == 0:
        bys_back_test_eval_score_dict = params["bys_eval_score_dict"].copy()

    return bys_back_test_eval_score_dict


def Batch_Bayesian_Search(params: dict, bys, ignore_params=[]):
    init_env(params)

    all_stime = datetime.now()
    log(f'开始批量贝叶斯炼丹，总轮数 {params["study_batch_num"]}', send_wx=True)
    dbg(f"参数: {params}")
    a_df = pd.DataFrame()
    work_path = os.path.join(
        params["data_path"], func.replace_macro_path(params, params["batch_work_path"])
    )
    study_batch_csv = os.path.join(
        work_path, func.replace_macro_path(params, params["study_batch_csv"])
    )
    if os.path.isfile(study_batch_csv):
        a_df = pd.read_csv(study_batch_csv, encoding="gbk")
    if a_df.empty:
        batch_id = datetime.now().strftime("%Y%m%d$%H%M%S")
        max_score = -30000
    else:
        batch_id = a_df["批次"].iloc[0]
        max_score = a_df["score"].max()

    # 取得用于计算评分的参数
    cls = __import__(
        f'program.{params["base_stg_folder"]}.{params["base_stg_file"]}', fromlist=("",)
    )
    stg_params = {}
    for key in params:
        if (
            not key in ignore_params
            and hasattr(cls, key)
            and getattr(cls, key) is not None
        ):
            params[key] = getattr(cls, key)
            stg_params[key] = params[key]
    log(f"从策略文件中更新配置参数: {str(stg_params)}")

    # 取得用于计算样本外回测评分的参数，如果没有配置则使用炼丹的参数
    bys_back_test_eval_score_dict = get_back_test_eval_score_dict(params)

    for n in range(1, params["study_batch_num"] + 1):
        if not a_df.empty and len(a_df[a_df["轮次"].astype(int) == n]) > 0:
            log(f"第 {n} 次炼丹结果已存在，跳过")
            continue

        stime = datetime.now()
        log("===========================================")
        log(f"开始第 {n} 次炼丹...", send_wx=True)
        log("===========================================")

        # 执行一整轮炼丹
        params["批次"] = batch_id
        params["轮次"] = f"{n:03}"
        params["disable_output_stg"] = True
        # 批次第一次启动标志
        if not "batch_first" in params.keys():
            params["batch_first"] = True
        else:
            params["batch_first"] = False
        # 最后一轮标志
        params["batch_last"] = n == params["study_batch_num"]

        result_df = Bayesian_Search(params.copy(), bys)

        # 对结果按 score 排序取最高的几个
        result_df.sort_values(
            by=["score"], ascending=[False], ignore_index=True, inplace=True
        )
        result_df.drop_duplicates(
            subset=["回测数据"], keep="first", ignore_index=True, inplace=True
        )
        df = result_df.head(params["output_stg_top_num_batch"])
        df["批次"] = batch_id
        df["轮次"] = n

        inner_score = df["score"].max()

        # 取得样本外回测的时间区间
        back_test_date_start = params["back_test_date_start"]
        back_test_date_end = params["back_test_date_end"]
        if back_test_date_start is None:
            back_test_date_start = params["sample_date_end"]

        if params["use_back_test_mode"] == False:
            # 样本外回测关闭
            log(f"样本外回测评分功能已禁用，最终的综合评分使用样本内评分")
        elif back_test_date_start is None:
            # 样本外起始不能为 None
            log(f"样本外回测起始日期为 None，禁用样本外回测功能")
        elif pd.to_datetime(
            back_test_date_end if back_test_date_end is not None else datetime.today()
        ) - pd.to_datetime(back_test_date_start) < pd.Timedelta(days=30):
            # 样本外回测时间间隔不能小于30天
            log(f"样本外回测时间间隔小于30天，禁用样本外回测功能")
        else:
            # 生成临时策略文件用于计算样本外表现，生成最终的评分
            df = run_top_stg_back_test(
                params,
                bys,
                df,
                cls,
                "样本外",
                f"轮次 {n}",
                back_test_date_start,
                back_test_date_end,
                bys_back_test_eval_score_dict,
                update_score=True,
            )

        n_score = df["score"].max()
        if n_score > max_score:
            msg = f"第 {n} 次炼丹完成，最大分数为 {n_score:.2f}，超过历史最高分数 {max_score:.2f}，用时 {datetime.now() - stime}"
            max_score = n_score
        else:
            msg = f"第 {n} 次炼丹完成，最大分数为 {n_score:.2f}，低于历史最高分数 {max_score:.2f}，用时 {datetime.now() - stime}"
        log(msg, send_wx=True)

        # 合并数据并保存
        df = df.sort_values("score", ascending=False, ignore_index=True)
        a_df = pd.concat([a_df, df.head(params["output_stg_top_num"])], ignore_index=True)
        a_df.sort_values(by=["score"], ascending=[False], inplace=True)
        a_df.to_csv(study_batch_csv, index=False, encoding="gbk")

        # 每轮测试完成都输出结果
        # 只有最后一轮并且配置允许时才打开回测图
        open_pic = (n == params["study_batch_num"]) and (
            not "open_pic" in params or params["open_pic"]
        )
        # 根据回测模式选择回测数量
        back_test_num = 0
        if (params["output_stg_batch_back_test_mode"] == "all") or (
            (params["output_stg_batch_back_test_mode"] == "last")
            and (n == params["study_batch_num"])
        ):
            back_test_num = params["output_stg_batch_back_test_num"]
        stg_file_list = bys.output_result_stg_files(
            params,
            a_df,
            params["output_stg_top_num_batch"],
            back_test_num,
            f"批量${batch_id}",
            open_pic=open_pic,
        )

        # 复制一份输出的策略文件存档
        round_path = os.path.join(
            params["data_path"],
            func.replace_macro_path(params, params["batch_finished_path"]),
        )
        finished_path = os.path.dirname(round_path)
        out_stg_path = os.path.join(finished_path, "结果策略")
        os.makedirs(out_stg_path, exist_ok=True)
        for stg_file in stg_file_list:
            fn = os.path.join(
                Cfg.root_path, "program", params["output_stg_folder"], f"{stg_file}.py"
            )
            if os.path.isfile(fn):
                shutil.copyfile(fn, os.path.join(out_stg_path, os.path.basename(fn)))

        # 更新内部数据索引文件
        org_file = params["study_storage_file"]
        arch_file = params["study_storage_file"].replace(".", "_arch.")
        shutil.move(
            os.path.join(round_path, org_file), os.path.join(round_path, arch_file)
        )
        storage_csv = os.path.join(
            finished_path,
            params["study_storage_file"].replace(".pkl", "") + "_批量.csv",
        )
        st_df = pd.DataFrame(
            [
                {
                    "炼丹策略": params["base_stg_file"],
                    "批次": batch_id,
                    "轮次": n,
                    "炼丹得分": inner_score,
                    "综合得分": n_score,
                    "内部数据": os.path.join(os.path.basename(round_path), arch_file),
                }
            ]
        )
        if os.path.isfile(storage_csv):
            st_df.to_csv(
                storage_csv, mode="a", encoding="gbk", header=False, index=False
            )
        else:
            st_df.to_csv(storage_csv, mode="w", encoding="gbk", index=False)

    # 取得批次目录
    finished_path = os.path.dirname(
        os.path.join(
            params["data_path"],
            func.replace_macro_path(params, params["batch_finished_path"]),
        )
    )
    shutil.move(
        study_batch_csv, os.path.join(finished_path, os.path.basename(study_batch_csv))
    )

    log(
        f'所有 {params["study_batch_num"]} 次炼丹完成，最高分 {max_score:.2f}，总用时 {datetime.now() - all_stime}',
        send_wx=True,
    )
    log(
        (
            f'最高分出现在第 {a_df["轮次"].iloc[0]} 轮次，迭代次数：{a_df["迭代"].iloc[0]}\n过滤参数：\n'
            + str(a_df["过滤参数"].iloc[0])
            if "过滤参数" in a_df.columns
            else (
                "" + "\n排序参数：\n" + str(a_df["排序参数"].iloc[0])
                if "排序参数" in a_df.columns
                else "" + "\n回测数据：\n" + str(a_df["回测数据"].iloc[0])
            )
        ),
        send_wx=True,
    )

    return a_df
